package cn.zwk.common.utils.tree;

import cn.zwk.common.model.TreeVo;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class TreeDrive<T extends TreeVo<T>> {
    private final List<T> list;
    private final Map<Long, TreeVo<T>> indexMap = new HashMap<>();
    private final TreeVo<T> root = new TreeRoot<>();
    private boolean invoke = false;
    private static final Long ROOT_ID = 0L;

    public TreeDrive(List<T> list) {
        this.list = list;
    }

    /*初始化树*/
    private void init() {
        if (!invoke) {
            buildTree();
            invoke = true;
        }
    }

    private void buildTree() {
        list.forEach(treeVo -> {
            indexMap.put(treeVo.getId(), treeVo);
        });
        indexMap.put(ROOT_ID, root);
        list.forEach(treeVo -> {
            TreeVo<T> parent = indexMap.get(treeVo.getParentId());
            if (parent != null) {
                parent.getChildren().add(treeVo);
            }
        });
        indexMap.remove(ROOT_ID);
    }

    public List<T> getTreeList() {
        init();
        return (List<T>) root.getChildren();
    }

    public Map<Long, TreeVo<T>> getIndexMap() {
        init();
        return indexMap;
    }

    public boolean hasLoop(Long id, Long pid) {
        init();
        while (pid != ROOT_ID) {
            if (pid.equals(id)) {
                return true;
            }
            pid = indexMap.get(pid).getParentId();
        }
        return false;
    }

//    public List<T> allChildren(Long id) {
//        init();
//        List<T> result = new ArrayList<>();
//        Queue<TreeVo> queue = new ArrayDeque<>();
//        queue.add(indexMap.get(id));
//        while (!queue.isEmpty()) {
//            TreeVo poll = queue.poll();
//            result.add((T) poll);
//            queue.addAll(poll.getChildren());
//        }
//        return result;
//    }

    /*获取该树的所有节点*/
    public List<T> getTreeNodes(Long id) {
        init();
        List<T> result = new ArrayList<>();
        result.add((T) indexMap.get(id));
        int head = 0, tail = 1;
        while (head < tail) {
            TreeVo<T> poll = result.get(head++);
            result.addAll(poll.getChildren());
            tail += poll.getChildren().size();
        }
        return result;
    }

    static class TreeRoot<T> extends TreeVo<T> {
        @Override
        public Long getId() {
            return ROOT_ID;
        }
    }
}
